#!/usr/bin/python
# -*-coding:utf-8-*-
import os
import pandas as pd
import numpy as np
from scipy import stats
import statsmodels.api as sm

from zg_data_process.zg_data_process import DataProcess

try:
    try:
        from zg_factor_analysis_module.base.FactorAnalysisLibDerivative import FactorAnalysisLibDerivative
        from zg_factor_analysis_module.base.dir_info import save_dir
    except:
        from .base.FactorAnalysisLibDerivative import FactorAnalysisLibDerivative
        from .base.dir_info import save_dir
except:
    from base.FactorAnalysisLibDerivative import FactorAnalysisLib_derivative
    from base.dir_info import save_dir


# TODO 回归因子收益率计算(日）
def cal_reg_factor_return(df, factor_name, ind='sw1', method='OLS'):
    '''
    :param df: 某一天的横截面数据，包括：因子值、总市值、流通市值、行业标签、个股收益率等数据项
    :param ind: 选择的行业分类标签 , 可选 sw1(申万一级行业分类)、ci1(中信一级行业分类)
    :param method: 计算因子收益率的方式，默认为OLS，可选WLS
    :return: WLS回归的因子收益率、t统计量 与p值
    '''
    reg_result = {}
    reg_result['factor_return'] = np.nan
    reg_result['tvalues'] = np.nan
    reg_result['pvalues'] = np.nan

    try:
        data = df.copy()
        # TODO 获取行业标签
        if ind == 'sw1':
            ind_data = pd.get_dummies(data['sw1_name'])
        else:
            ind_data = pd.get_dummies(data['ci1_name'])

        data[ind_data.columns] = ind_data
        reg_data = data[ind_data.columns.tolist() +
                        [factor_name, 'scale_circulate_market_size', 'scale_total_market_size', 'quote_rate']]
        # TODO 加权最小二乘回归法计算因子收益率，权重为流通市值的平方根
        if method == 'WLS':
            reg = sm.WLS(reg_data['quote_rate'],
                         sm.add_constant(reg_data[['scale_total_market_size', factor_name] + ind_data.columns.tolist()[1:]]),
                         weights=np.sqrt(reg_data['scale_circulate_market_size'])).fit()
        else:
            reg = sm.OLS(reg_data['quote_rate'],
                         sm.add_constant(reg_data[['scale_total_market_size', factor_name] + ind_data.columns.tolist()[1:]])
                         ).fit()
        reg_result['factor_return'] = reg.params[factor_name]
        reg_result['tvalues'] = reg.tvalues[factor_name]
        reg_result['pvalues'] = reg.pvalues[factor_name]
    except Exception as e:
        print(e)
        pass

    return reg_result

# TODO 区间回归统计量计算
def cal_reg_stats(reg_result):
    reg_stats = {}

    reg_stats['factor_return_mean'] = reg_result['factor_return'].mean()

    reg_stats['factor_return_std'] = reg_result['factor_return'].std()

    reg_stats['tvalues_mean'] = reg_result['tvalues'].mean()

    reg_stats['pvalues_mean'] = reg_result['pvalues'].mean()

    reg_stats['abs_tvalues_mean'] = reg_result['tvalues'].abs().mean()

    reg_stats['abs_tvalues>2_ratio'] = \
        reg_result[reg_result['tvalues'].abs() >= 2].shape[0] / reg_result.shape[0]

    reg_stats['Fama_MacBeth_tvalues'] = \
        reg_result['factor_return'].mean() / reg_result['factor_return'].std() * np.sqrt(reg_result.shape[0] - 1)

    return reg_stats

# TODO 计算因子回归收益率
def factor_reg_return(concat_data, factor_name, ind='sw1', method='OLS'):
    """
    :param concat_data:  concat_data 为经过处理后的合并数据，包括因子（多个）、行业、总市值、流通市值、个股收益率等数据
    :return: 每天的因子收益率和T统计量与P值
    """
    cal_data = concat_data.set_index(['date', 'stock_code']).copy()
    reg_items = [factor_name, 'scale_circulate_market_size', 'scale_total_market_size', 'sw1_name']

    for item in reg_items:
        # print(item)
        cal_data[item] = cal_data.groupby('stock_code')[item].shift(1)

    cal_data = cal_data.dropna()
    total_result = {}

    for index, group in cal_data.groupby('date'):
        # print(index)
        total_result[index] = cal_reg_factor_return(group, factor_name, ind, method)

    reg_result = pd.DataFrame(total_result).T
    reg_stats = cal_reg_stats(reg_result)

    return reg_result, reg_stats

# TODO 计算因子IC
def cal_factor_ic(concat_data, factor_name, periods=21, method='rank'):
    cal_data = concat_data.set_index(['date', 'stock_code']).copy()
    mr = cal_data.groupby('code')['quote_rate'].rolling(periods).sum()
    mr.index = mr.index.droplevel('code')
    cal_data['monthly_return'] = mr
    cal_data[factor_name] = cal_data.groupby('stock_code')[factor_name].shift(periods + 1)
    cal_data = cal_data.dropna()

    # TODO 计算加权IC
    def weight_ic(data):
        temp = data.copy()
        temp['weight'] = temp[factor_name].rank().apply(lambda x: 2 ** (2 * x / temp.shape[0] - 2))
        temp['weight_x'] = temp[factor_name] * temp['weight']
        temp['weight_r'] = temp['monthly_return'] * temp['weight']
        return stats.spearmanr(temp['weight_x'], temp['weight_r'])[0]

    if method == 'rank':
        ic = cal_data.groupby('date').apply(lambda x: stats.spearmanr(x[factor_name], x['monthly_return'])[0])
    elif method == 'weight':
        ic = cal_data.groupby('date').apply(lambda x: weight_ic(x))
    else:
        ic = cal_data.groupby('date').apply(lambda x: stats.pearsonr(x[factor_name], x['monthly_return'])[0])
    ic.name = factor_name

    # TODO 计算IC统计量
    def cal_ic_statis(ic):
        ic_stats = {}
        ic_stats['ic_mean'] = ic.dropna().mean()
        ic_stats['ic_std'] = ic.dropna().std()
        ic_stats['ic_ir'] = ic_stats['ic_mean'] / ic_stats['ic_std']
        ic_stats['pos_ic_ratio'] = ic[ic > 0].shape[0] / ic.shape[0]
        return ic_stats

    ic_stats = cal_ic_statis(ic)
    return ic, ic_stats


# TODO 添加分组标签
def add_group_lebel(df, factor_name, method='cap', quantile=5, group_num=5):
    df = df.reset_index('date').copy()
    if method == 'cap':
        grouping = pd.qcut(df['scale_total_market_size'].rank(method='first'), group_num, labels=False)
        grouped = df[factor_name].rank(method='first').groupby(grouping).apply(pd.qcut, quantile,
                                                                               labels=False) + 1
    elif method == 'sw1':
        grouping = df['sw1_name']
        grouped = df[factor_name].rank(method='first').groupby(grouping).apply(pd.qcut, quantile,
                                                                               labels=False) + 1
    else:
        grouped = pd.qcut(df[factor_name].rank(method='first'), quantile, labels=False) + 1

    return grouped

# TODO 计算股票权重
def cal_stock_weight(df, method='ew'):
    if method == 'cir_cap':
        weight = df.groupby('group').apply(
            lambda x: x['scale_circulate_market_size'] / x['scale_circulate_market_size'].sum())
    elif method == 'total_cap':
        weight = df.groupby('group').apply(
            lambda x: x['scale_total_market_size'] / x['scale_total_market_size'].sum())
    else:
        df['weight'] = 1
        weight = df.groupby('group').apply(lambda x: x['weight'] / x.shape[0])

    weight.index = weight.index.droplevel('group')
    weight = weight.sort_index()

    return weight

# TODO data filter
def fillter_analysis_data(concat_data, filter_st=True, filter_sub_new=True, filter_paused=True, pool=None):
    data_process = DataProcess()
    concat_data = concat_data.set_index(['stock_code', 'date'])

    # ST过滤
    if filter_st:
        concat_data = data_process.st_filteration(concat_data, drop_label=False)
    # 停牌股过滤
    if filter_paused:
        concat_data = data_process.paused_stock_filtration(concat_data, drop_label=False)
    # 新股次新股过滤
    if filter_sub_new:
        concat_data = data_process.new_subnew_stock_filtration(concat_data, drop_label=False)
    # 股票池过滤
    if pool == '000905':
        concat_data = data_process.get_csi500_component_data(concat_data, drop_label=False)
    elif pool == '000300':
        concat_data = data_process.get_csi300_component_data(concat_data, drop_label=False)
    else:
        pass

    print('filter analysis data finish!')

    return concat_data.reset_index()

# TODO - 获取指数涨跌幅数据
def get_daily_index_trade_data(benchmark):
    index_trade_data = data_reader.read_basic_data_table('processed_daily_index_trade_data')

    index_trade_data = index_trade_data[index_trade_data['index_code'] == benchmark]

    index_trade_data['quote_rate'] = index_trade_data['close'].pct_change()

    index_trade_data.index = pd.to_datetime(index_trade_data.date)

    return index_trade_data


# TODO 计算因子分组收益率
def cal_quantile_return(concat_data,
                        factor_name,
                        pool=None,
                        fillter=True,
                        bench='000905',
                        rebalance=21,
                        quantile=5,
                        group='cap',
                        weight='ew',
                        save_dir=save_dir):
    # TODO 数据过滤,若caocat_data已经过滤则不需要这一步
    if fillter:
        concat_data = fillter_analysis_data(concat_data, pool=pool)

    cal_data = concat_data.set_index(['date', 'stock_code']).copy()
    trade_date = sorted(cal_data.index.get_level_values('date').unique().tolist())
    rebalance_date = trade_date[::rebalance]

    # TODO 添加因子分组标签
    grouped_data = cal_data.groupby('date').apply(
        lambda x: add_group_lebel(x, factor_name, method=group, quantile=quantile))
    cal_data['group'] = grouped_data
    fill_date = sorted(list(set(trade_date).difference(set(rebalance_date))))
    cal_data.loc[fill_date, 'group'] = np.nan
    cal_data['group'] = cal_data.groupby('stock_code')['group'].fillna(method='ffill')
    cal_data['group'] = cal_data.groupby('stock_code')['group'].shift(1)

    # TODO 计算个股权重
    cal_data['dt'] = cal_data.index.get_level_values('date')
    w_data = cal_data.dropna()
    weights = w_data.groupby('dt').apply(lambda x: cal_stock_weight(x, method=weight))
    weights.index = weights.index.droplevel('dt')
    cal_data['weight'] = weights

    # TODO 计算加权收益率
    cal_data['weight_return'] = cal_data['weight'] * cal_data['quote_rate']
    # cal_data['weight_excess_return'] = cal_data['weight'] * cal_data['excess_return']

    # TODO 计算分组收益率
    ret = cal_data.groupby(['date', 'group'])['weight_return'].sum()
    ret = ret.reset_index().pivot(index='date', columns='group', values='weight_return')
    ret['top-bot'] = ret[5] - ret[1]
    index_ret = get_daily_index_trade_data(bench)
    ret['bench'] = index_ret.loc[ret.index, 'quote_rate']
    cum_ret = (1 + ret).cumprod()
    excess_ret = ret.iloc[:, 0:quantile].apply(lambda x: x - ret['bench'])
    cum_excess_ret = (1 + excess_ret).cumprod()

    # TODO 获取调仓记录数据
    cal_data[['group', 'weight']] = cal_data[['group', 'weight']].fillna(method='bfill')
    position_data = cal_data.reset_index('stock_code').loc[rebalance_date].reset_index()
    writer = pd.ExcelWriter(os.path.join(save_dir, factor_name + '_factor_groupby_' + str(group) + '_' + str(
        weight) + '_weight_' + pool + '_pool_' + str(rebalance) + '_quantile_test_result.xlsx'))
    ret.to_excel(writer, 'abs_return_result')
    cum_ret.to_excel(writer, 'cum_abs_return_result')
    excess_ret.to_excel(writer, 'excess_return_result')
    cum_excess_ret.to_excel(writer, 'cum_excess_return_result')
    position_data.to_excel(writer, 'holding_detail')
    writer.save()

    def cal_annual_return(daily_return):
        cum_return = (1 + daily_return).cumprod()
        annual_return = cum_return.ix[-1] ** (243 * 1. / len(cum_return)) - 1
        return annual_return

    def cal_annual_vol(daily_return):
        return daily_return.std() * np.sqrt(243)

    def cal_sharpe(daily_return):
        annual_return = cal_annual_return(daily_return)
        annual_vol = cal_annual_vol(daily_return)
        sharpe = annual_return / annual_vol
        return sharpe

    def cal_winrate(daily_return):
        win_rate = len(daily_return[daily_return > 0]) / len(daily_return)

        return win_rate

    def cal_max_draw(daily_return):
        cum_return = (1 + daily_return).cumprod()
        cum_peak = np.maximum.accumulate(cum_return)
        dd_to_cum_peak = (cum_peak - cum_return) / cum_peak
        max_dd_end = np.argmax(dd_to_cum_peak)
        max_dd = dd_to_cum_peak[max_dd_end]

        return max_dd

    quantile_result_stats = {}
    quantile_result_stats['annual_return'] = ret.apply(lambda x: cal_annual_return(x))
    quantile_result_stats['annual_vol'] = ret.apply(lambda x: cal_annual_vol(x))
    quantile_result_stats['SR'] = ret.apply(lambda x: cal_sharpe(x))
    quantile_result_stats['winrate'] = ret.apply(lambda x: cal_winrate(x))
    quantile_result_stats['max_draw'] = ret.apply(lambda x: cal_max_draw(x))
    quantile_result_stats['excess_annual_return'] = excess_ret.apply(lambda x: cal_annual_return(x))
    quantile_result_stats['TE'] = excess_ret.apply(lambda x: cal_annual_vol(x))
    quantile_result_stats['IR'] = excess_ret.apply(lambda x: cal_sharpe(x))
    quantile_result_stats['excess_winrate'] = excess_ret.apply(lambda x: cal_winrate(x))
    quantile_result_stats['excess_max_draw'] = excess_ret.apply(lambda x: cal_max_draw(x))

    return ret, excess_ret, quantile_result_stats

# TODO 获取因子持仓数据
def get_factor_position_data(factor_name, pool=None, rebalance=21, group='cap', weight='ew'):
    data = pd.read_excel(os.path.join(save_dir,
                                      factor_name + '_factor_groupby_' +
                                      str(group) + '_' +
                                      str(weight) + '_weight_' +
                                      pool + '_pool_' +
                                      str(rebalance) + '_quantile_test_result.xlsx'),
                         'holding_detail')

    data['group'] = data['group'].astype('str')

    data = data.set_index(['date', 'group'])

    return data

# TODO 获取因子回归收益数据
def get_factor_reg_return(factor_name,pool='000905'):
    data = pd.read_excel(os.path.join(save_dir, factor_name + '_'+pool+'_pool_analysis_result.xlsx'), 'reg_result')

    return data

# TODO 获取因子IC数据
def get_factor_ic_data(factor_name,pool='000905'):
    data=pd.read_excel(os.path.join(save_dir, factor_name + '_'+pool+'_pool_analysis_result.xlsx'), 'ic_result')

    return data

#TODO 计算因子的选股交叉度——在分层回测的持仓数据基础之上
def cal_factors_stock_overlap(factor1, factor2, pool=None, rebalance=21, group='cap', weight='ew'):
    try:
        p1 = get_factor_position_data(factor1, pool=pool, rebalance=rebalance, group=group, weight=weight)
        p2 = get_factor_position_data(factor2, pool=pool, rebalance=rebalance, group=group, weight=weight)
        com_date_list = sorted(
            list(p1.index.get_level_values('date').unique().intersection(p2.index.get_level_values('date').unique())))
        temp = {}

        for date in com_date_list:
            df1 = p1.xs(date)
            df2 = p2.xs(date)
            overlap_ratio = {}
            for q in df1.index.unique():
                stocks1 = df1.loc[q]['stock_code'].tolist()
                stocks2 = df2.loc[q]['stock_code'].tolist()
                overlap_ratio[q] = (len(set(stocks1).intersection(set(stocks2))) / max(len(stocks1), len(stocks2)))
            overlap_ratio['mean'] = pd.Series(overlap_ratio).mean()
            temp[date] = overlap_ratio

        return pd.DataFrame(temp).T
    except Exception as e:
        print(e)
        return np.nan

#TODO 计算因子回归收益率的相关性——在计算出回归因子收益率的基础上
def cal_reg_factor_return_corr(factor1,factor2,pool='000905'):
    df1=get_factor_reg_return(factor1,pool)

    df2=get_factor_reg_return(factor2,pool)

    df=pd.concat([df1['factor_return'],df2['factor_return']],axis=1)

    corr=df.corr().iloc[0,1]

    return corr

#TODO 计算因子IC的相关性——在计算出IC的基础上
def cal_factor_ic_corr(factor1, factor2, pool='000905'):
    df1=get_factor_ic_data(factor1,pool)

    df2=get_factor_ic_data(factor2,pool)

    df=pd.concat([df1[factor1],df2[factor2]],axis=1)

    corr=df.corr().iloc[0,1]

    return corr


if __name__ == '__main__':
    data_reader = FactorAnalysisLibDerivative()

    factor='growth_roic_yoy'

    concat_data=data_reader.read_zbc_factor_table('concat_test_data')

    #TODO 因子分层回测
    quantile_return,\
    excess_quantile_return,\
    factor_return_stats =\
        cal_quantile_return(concat_data,
                            factor,
                            pool='000905',
                            fillter=True,
                            bench='000905',
                            rebalance=21,
                            quantile=5,
                            group='cap',
                            weight='ew',
                            save_dir=save_dir)

    # TODO 因子IC计算
    ic, ic_stats = cal_factor_ic(concat_data, factor, periods=21, method='rank')

    #TODO 因子回归收益计算
    reg_return, reg_stats = factor_reg_return(concat_data, factor, ind='sw1', method='OLS')

    #TODO 计算两个因子的选股重叠度
    factor1='growth_roic_yoy'
    factor2='growth_marginal_roic'
    stocks_overlap=cal_factors_stock_overlap(factor1, factor2, pool='000905', rebalance=21, group='cap', weight='ew')

    #TODO 计算两个因子的IC相关性
    ic_corr=cal_factor_ic_corr(factor1,factor2,pool='000905')

    #TODO 计算两个因子的回归因子收益率相关性
    reg_corr =cal_reg_factor_return_corr(factor1, factor2, pool='000905')
